from sqlalchemy import create_engine, Column, Date
from sqlalchemy.sql import text
from datetime import datetime
import sqlite3


def migrate_database():
    # 连接到现有数据库
    conn = sqlite3.connect("predictions.db")
    cursor = conn.cursor()

    try:
        # 1. 创建临时表
        cursor.execute("""
            CREATE TABLE prediction_records_temp (
                id INTEGER PRIMARY KEY,
                timestamp DATETIME,
                predicted_value FLOAT,
                date DATE
            )
        """)

        # 2. 从原表复制数据到临时表，同时添加date字段
        cursor.execute("""
            INSERT INTO prediction_records_temp (id, timestamp, predicted_value, date)
            SELECT id, timestamp, predicted_value, date(timestamp)
            FROM prediction_records
        """)

        # 3. 删除原表
        cursor.execute("DROP TABLE prediction_records")

        # 4. 将临时表重命名为原表名
        cursor.execute(
            "ALTER TABLE prediction_records_temp RENAME TO prediction_records"
        )

        # 5. 提交更改
        conn.commit()
        print("数据库迁移成功！")

    except Exception as e:
        conn.rollback()
        print(f"迁移失败: {str(e)}")
        raise
    finally:
        conn.close()


if __name__ == "__main__":
    migrate_database()
